# coding=utf-8
### 生产者
import json
import socket
import pika
import time

hostname = socket.gethostname()
ip_address = socket.gethostbyname(hostname)

user_info = pika.PlainCredentials('admin', 'admin')
connection = pika.BlockingConnection(
    pika.ConnectionParameters('192.168.1.111', 5672, virtual_host='/', credentials=user_info))
channel = connection.channel()
queue_name = 'inference'
channel.queue_declare(queue=queue_name, durable=True)  # 持久化

# body = {
#     'type': 'inference',
#     'train_task_id': 1,
#     'inference_items': [
#         {'id': 1, 'url': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2024/12/29/0a26aceb2_index_3.png'},
#         {'id': 2, 'url': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2024/12/29/0a37f0a29_index_2.png'},
#         {'id': 3, 'url': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2024/12/29/0a8cb8ddf_index_1.png'},
#         {'id': 4, 'url': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2024/12/29/0a8cb8ddf_index_4.png'},
#     ],
#     'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/inference_config.yaml'
# }

body = config = {
    'type': 'deploy',
    'train_task_id': 7,
    'task_type': 'segment',
    'segment_model': {
        'name': 'resnet34',
        'pretrained': False,
        'class_num': 4,
        'parameters': {'activation': None},
        'tta_flag': False,  # 对图像翻转等操作，多种预测取均值
        'checkpoint_link': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2025/01/08/a8934b79-8249-4fdd-bb53-e85d2ac316c3.pth',  # 模型权重保存位置
        'local_path': './oss_model_path/segment',
    },
    'classify_model': {
        'name': 'resnet18',
        'pretrained': False,
        'class_num': 4,
        'parameters': {},
        'tta_flag': False,  # 对图像翻转等操作，多种预测取均值
        'checkpoint_link': 'http://msjava.oss-cn-qingdao.aliyuncs.com/2025/01/07/7f1a5724-57c6-41e1-bb31-1e7104c51178.pth',  # 模型权重保存位置
        'local_path': './oss_model_path/classify',
    },
    'dataset': {
        'mask_color': [[255, 0, 0], [0, 255, 0], [0, 0, 255], [139, 0, 139]],
        'height': 256,
        'width': 1600,
        'mean': [0.485, 0.456, 0.406],
        'std': [0.229, 0.224, 0.225],
    },
}

channel.basic_publish(exchange='',  # 当前是一个简单模式，所以这里设置为空字符串就可以了
                      routing_key=queue_name,  # 指定消息要发送到哪个queue
                      body=json.dumps(body)  # 指定要发送的消息
                      )
